import numpy as np
import matplotlib.pyplot as plt
from init_control import get_control_parameters
from pid_controller import PIDController, CascadePIDController
from uav_dynamics import uav_dynamics


def desired_euler_angles(eax, eay, max_angle_roll, max_angle_pitch):
    g = 9.8  # 重力加速度 (m/s^2)
    phi = eay / g
    theta = -eax / g

    # 饱和处理
    phi = np.clip(phi, -max_angle_roll * DEG2RAD, max_angle_roll * DEG2RAD)
    theta = np.clip(theta, -max_angle_pitch * DEG2RAD, max_angle_pitch * DEG2RAD)

    return phi, theta


def motor_mixer(roll, pitch, yaw, thrust):
    idle_PWM = 1000
    scale = 1000

    M1 = (thrust - roll + pitch + yaw) * scale + idle_PWM
    M2 = (thrust + roll - pitch + yaw) * scale + idle_PWM
    M3 = (thrust + roll + pitch - yaw) * scale + idle_PWM
    M4 = (thrust - roll - pitch - yaw) * scale + idle_PWM

    # 防止负值
    M1 = max(0, M1)
    M2 = max(0, M2)
    M3 = max(0, M3)
    M4 = max(0, M4)

    return M1, M2, M3, M4


def run_uav_simulation():
    params = get_control_parameters()

    # 初始化PID控制器
    roll_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0)
    roll_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
                                         params['Kd_RP_AngleRate'],
                                         integrator_max=params['Saturation_I_RP_Max'],
                                         integrator_min=params['Saturation_I_RP_Min'])
    roll_controller = CascadePIDController(roll_angle_controller, roll_rate_controller)

    pitch_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0)
    pitch_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
                                          params['Kd_RP_AngleRate'],
                                          integrator_max=params['Saturation_I_RP_Max'],
                                          integrator_min=params['Saturation_I_RP_Min'])
    pitch_controller = CascadePIDController(pitch_angle_controller, pitch_rate_controller)

    yaw_rate_controller = PIDController(params['Kp_YAW_AngleRate'], params['Ki_YAW_AngleRate'],
                                        params['Kd_YAW_AngleRate'],
                                        integrator_max=params['Saturation_I_Y_Max'],
                                        integrator_min=params['Saturation_I_Y_Min'])

    # 添加高度控制PID
    height_controller = PIDController(0.6, 0.3, 0.05)  # 调整这些值以获得稳定的高度控制

    # 初始化仿真状态
    position = params['ModelInit_PosE'].astype(np.float64)
    velocity = params['ModelInit_VelB'].astype(np.float64)
    angle = params['ModelInit_AngEuler'].astype(np.float64)
    rate = params['ModelInit_RateB'].astype(np.float64)
    w = np.zeros(4, dtype=np.float64)  # 电机转速

    dt = 0.01  # 时间步长，单位为秒
    simulation_time = 60  # 仿真总时间，单位为秒
    time_steps = int(simulation_time / dt)

    # 定义起飞阶段和正方形轨迹的四个角点
    target_altitude = -5  # 目标高度
    takeoff_time = 10  # 起飞时间，单位为秒
    takeoff_steps = int(takeoff_time / dt)

    square_corners = [
        np.array([0, 0, target_altitude]),  # 起点
        np.array([5, 0, target_altitude]),
        np.array([5, 5, target_altitude]),
        np.array([0, 5, target_altitude]),
        np.array([0, 0, target_altitude])  # 回到起点
    ]

    current_corner_index = 0
    next_corner_time = (time_steps - takeoff_steps) // len(square_corners)

    # 数据记录
    position_history = []

    for step in range(time_steps):
        if step < takeoff_steps:
            # 起飞阶段，垂直上升到目标高度
            setpoint = np.array([0, 0, target_altitude])
        else:
            # 正方形轨迹阶段
            total_steps_per_segment = (time_steps - takeoff_steps) // (len(square_corners) - 1)
            segment_start_index = current_corner_index % len(square_corners)
            segment_end_index = (current_corner_index + 1) % len(square_corners)
            segment_start_point = square_corners[segment_start_index]
            segment_end_point = square_corners[segment_end_index]
            progress = (step - takeoff_steps) % total_steps_per_segment / total_steps_per_segment
            setpoint = segment_start_point + progress * (segment_end_point - segment_start_point)

        # 计算位置误差
        eax = setpoint[0] - position[0]
        eay = setpoint[1] - position[1]

        # 期望的欧拉角
        phi, theta = desired_euler_angles(eax, eay, params['MAX_CONTROL_ANGLE_ROLL'], params['MAX_CONTROL_ANGLE_PITCH'])

        # 更新控制器输出
        roll_output = roll_controller.update(outer_setpoint=phi, outer_measurement=angle[0], inner_measurement=rate[0],
                                             dt=dt)
        pitch_output = pitch_controller.update(outer_setpoint=theta, outer_measurement=angle[1],
                                               inner_measurement=rate[1], dt=dt)
        yaw_output = yaw_rate_controller.update(setpoint=0, measurement=rate[2], dt=dt)

        # 高度控制
        altitude_output = height_controller.update(setpoint[2], position[2], dt)

        # 计算推力，考虑重力补偿
        thrust = (altitude_output + 1) * params['ModelParam_uavMass'] * params['ModelParam_envGravityAcc']

        # 计算电机转速
        M1, M2, M3, M4 = motor_mixer(roll_output, pitch_output, yaw_output, thrust)

        # 将PWM信号转换为转速（假设线性关系）
        w = np.array([pwm_to_rpm(M1), pwm_to_rpm(M2), pwm_to_rpm(M3), pwm_to_rpm(M4)])

        # 更新无人机动力学模型
        position, velocity, angle, rate = uav_dynamics(position, velocity, angle, rate, w, dt, params)

        # 记录位置
        position_history.append(position.copy())

        # 打印当前仿真状态（仅供示例）
        print(f"Time: {step * dt:.2f}s, Position: {position}, Velocity: {velocity}, Angle: {angle}, Rate: {rate}")

        # 更新当前角点索引
        if step >= takeoff_steps and (step - takeoff_steps) % total_steps_per_segment == 0:
            current_corner_index = (current_corner_index + 1) % (len(square_corners) - 1)

    # 转换位置历史记录为数组
    position_history = np.array(position_history)

    # 绘制无人机轨迹
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot(position_history[:, 0], position_history[:, 1], position_history[:, 2], label='UAV Trajectory')
    ax.scatter(square_corners[0][0], square_corners[0][1], square_corners[0][2], color='red', label='Start Point')
    ax.scatter(square_corners[-1][0], square_corners[-1][1], square_corners[-1][2], color='blue', label='End Point')
    ax.set_xlabel('X Position (m)')
    ax.set_ylabel('Y Position (m)')
    ax.set_zlabel('Z Position (m)')
    ax.legend()
    plt.show()


# 运行仿真
run_uav_simulation()
